package cn.johnyu;

import cn.johnyu.service.MilvusService;
import cn.johnyu.util.AliEmbeddingClient;
import cn.johnyu.util.MilvusUtil;
import cn.johnyu.util.QwenPlusClient;
import io.milvus.client.MilvusServiceClient;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

public class RagQwenOperate {

    public static void main(String[] args) throws Exception{
        String collectionName = "law_articles106";
        int dim=1024;
        String question="我使用爬虫技术盗取了某公司的敏感数据，并用于商用，应该触犯了哪条刑法";

        //获取问题的向量
        List<Float> queryVector = AliEmbeddingClient.embedding(Arrays.asList(question)).get(0);

        MilvusServiceClient milvusServiceClient = MilvusUtil.getClient("localhost", 19530);

        //向量搜索
        List<Map<String, String>> results =
                MilvusService.searchData(milvusServiceClient, collectionName, queryVector);

        String context = results.stream()
                .map(map -> map.get("text"))
                .reduce("", (a, b) -> a + "\n" + b);


        context="请根据以下法律条文内容回答用户问题，并适当结合一些中国相关法律信息进行回答：\n"+context;

        String rs= QwenPlusClient.callWithMessage(context, question);

        System.out.println("经过Rag后的回答:\n " + rs);

        MilvusUtil.close();
    }
}
